#include <iostream>
#include <climits>
#include <LEDA/graphics/graphwin.h>
#include <LEDA/graphics/color.h>
#include <LEDA/system/basic.h>
#include "control.h"

#include <LEDA/core/partition.h>
#include <LEDA/core/p_queue.h>
#include <LEDA/graph/node_partition.h>


using namespace leda;


int get_edge_weight(GraphWin &gw , edge &e) {
	string e_label = gw.get_user_label(e);
	int weight = atoi(e_label.c_str());
	return weight;
}


void kruskal(graph &g, GraphWin &gw, int &error) {

	int current_mst_weight = 0;

	p_queue<int, edge> priority_queue;
	node_partition sets(g);
	edge e;
	node v;

	//Put edges into priority queue (ordered by their edge weight)
	forall_edges(e, g) {
		int weight = get_edge_weight(gw , e);
		priority_queue.insert(weight, e);
	}

	//generate a partition for each node in the graph
	forall_nodes(v,g) {
		sets.make_rep(v);
	}

	//pop edges from the priority queue until it is empty or there is only 1 partition left.
	//if we pop an edge that connects 2 different partitions we union them.
	while (sets.number_of_blocks() > 1 && !priority_queue.empty()) {

		pq_item pq_edge = priority_queue.find_min();
		e = priority_queue.inf(pq_edge);

		node start;
		node target;

		//get first node of edge
		start = g.opposite(e, start);
		//get second node of edge
		target = g.opposite(e, start);


		//mark edge that is about to get checked out
		gw.set_color(e , green);
		gw.set_width(e, 2);
		control_wait(0.5);

		// If they do not belong to the same partition union the partitions and add the edge to the MST
		if (!sets.same_block(start , target))
		{
			current_mst_weight += get_edge_weight(gw,e);

			sets.union_blocks(start , target);
			gw.set_color(e , blue);
			gw.set_width(e, 3);
			gw.set_color(start , blue);
			gw.set_color(target , blue);
			control_wait(0.5);

		}
		// Else both partitions are the same and we can not use this edge. In this case we mark it as visited but unused
		else
		{
			gw.set_width(e, 2);
			gw.set_color(e,red);
		}
		//delete item from the priority queue
		priority_queue.del_item(pq_edge);
	}

	//If we could not union all partitions the graph can not be connected
	if (sets.number_of_blocks() > 1)
	{
		error = 1;
	}


	gw.message("MST weight: " + string("%d" , current_mst_weight));

}


int main(int argc, char *argv[]) {

    GraphWin gw(800, 600);
    gw.display();
    create_control();
    gw.set_directed(false);
    if (argc > 1)
        gw.read(argv[1]);
    gw.edit();
    graph &g = gw.get_graph();
    if (g.number_of_nodes() == 0) {
        gw.close(); destroy_control();
        exit(1);
    }

    node v;
    edge e;
    int error = 0;

    // initialize graph nodes and edges as yellow
    forall_nodes(v, g) {
        gw.set_label_type(v, user_label);
        gw.set_color(v, grey3);
    }
    forall_edges(e, g)
    {
    	gw.set_color(e, grey3);
    }

    //run kruskal algorithm
    kruskal(g , gw , error);

    if (error == 0) {
    	//tell user that everything went well.
    	gw.acknowledge("Ready!");
    } else {
    	//tell user that there is a problem.
    	gw.acknowledge("The graph does not seem to be connected.");
    }


    gw.edit();
    gw.close();
    destroy_control();
    exit(0);

	return 0;


}
